"""从匹配好的三维点和二维点对解算相机位姿"""

# from g2o import g2opy as g2o
import numpy as np
from math import sin, cos
import cv2
from scipy.spatial.transform import Rotation
import os
import argparse
from g2o import g2opy as g2o


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "-c",
        "--collect_dir",
        type=str,
        help="The directory of the collected data",
    )

    return arg_parser.parse_args()


def Rotate(ori):
    """
    三轴转角
    """
    x, y, z = ori[2], ori[0], ori[1]
    R_x = np.array(
        [
            [1, 0, 0],
            [0, cos(x), sin(x)],
            [0, -sin(x), cos(x)],
        ]
    )
    R_y = np.array(
        [
            [cos(y), 0, -sin(y)],
            [0, 1, 0],
            [sin(y), 0, cos(y)],
        ]
    )
    R_z = np.array(
        [
            [cos(z), sin(z), 0],
            [-sin(z), cos(z), 0],
            [0, 0, 1],
        ]
    )

    return R_x @ R_y @ R_z


def real_pose(pose):
    # R = Rotate(pose[3:])
    R = Rotation.from_quat(pose[3:]).as_matrix().T
    T = -R @ np.array(pose[:3])
    R_ = np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
    H = np.hstack((R, T[:, None]))
    H = np.vstack((H, np.array([0, 0, 0, 1])))
    return R_ @ H


def read_points(file_path):
    uv = []
    XYZ = []
    with open(file_path, "r") as f:
        f.readline()
        for line in f:
            line = line.strip().split(",")
            x, y, x0, y0, z0 = [float(x) for x in line[:-1]]
            uv.append([x, y])
            XYZ.append([x0, y0, z0])
    return uv, XYZ


def initial_pose(K, observations, points_3d, use_ransac=True, th=8, n=4):
    def pose_calculate(X, x):
        distCoeffs = np.zeros(4)
        # 计算旋转向量和平移向量
        success, rvec, tvec = cv2.solvePnP(X, x, K, distCoeffs)
        if not success:
            Warning("Failed to solve PnP")
            return None, None
        # 将旋转向量转换为旋转矩阵
        R, _ = cv2.Rodrigues(rvec)
        return R, tvec

    if not use_ransac:
        return pose_calculate(points_3d, observations)
    else:
        best_num = 0
        best_inliers = np.array([False] * len(observations))
        sample = np.random.choice(len(observations), n, replace=False)
        for _ in range(100):
            R, tvec = pose_calculate(points_3d[sample], observations[sample])
            if R is not None:
                err = np.linalg.norm(
                    cv2.projectPoints(points_3d, R, tvec, K, np.zeros(4))[0].squeeze()
                    - observations,
                    axis=1,
                )
                inliers = err < th
                if inliers.sum() > best_num:
                    best_num = inliers.sum()
                    best_inliers = inliers
                if inliers.sum() > 0.8 * len(observations):
                    break
                # 另一种方式，使用绝对点个数，实际上超过20个点后计算所得精度已经相当高了。
                # if inliers.sum() > 20:
                #     break
                if inliers.sum() > n:
                    sample = np.random.choice(np.where(inliers)[0], n, replace=False)
                else:
                    sample = np.random.choice(len(observations), n, replace=False)
            else:
                sample = np.random.choice(len(observations), n, replace=False)
        if best_inliers.sum() < n:
            return None, None
        else:
            return pose_calculate(points_3d[best_inliers], observations[best_inliers])


def bundle_adjustment(K, observations, points_3d, initial_pose):
    # 初始化优化器
    optimizer = g2o.SparseOptimizer()
    solver = g2o.BlockSolverSE3(g2o.LinearSolverDenseSE3())
    solver = g2o.OptimizationAlgorithmLevenberg(solver)
    optimizer.set_algorithm(solver)

    # 添加相机参数顶点
    camera_params = g2o.CameraParameters(K[0, 0], K[:2, 2], 0)
    camera_params.set_id(0)
    optimizer.add_parameter(camera_params)
    # 假设有M个三维点
    M = len(points_3d)  # 三维点数量
    # 添加相机顶点
    v_se3 = g2o.VertexSE3Expmap()
    v_se3.set_id(0)
    v_se3.set_estimate(g2o.SE3Quat(*initial_pose))
    optimizer.add_vertex(v_se3)

    # 添加三维点顶点
    for j in range(M):
        v_p = g2o.VertexPointXYZ()
        v_p.set_id(1 + j)
        v_p.set_marginalized(True)
        v_p.set_estimate(points_3d[j])
        v_p.set_fixed(True)
        optimizer.add_vertex(v_p)

    # 添加边：三维点向二维点的投影
    for j in range(M):
        if observations[j] is not None:  # 检查是否有观测
            edge = g2o.EdgeProjectXYZ2UV()
            edge.set_vertex(0, optimizer.vertex(1 + j))  # 三维点顶点
            edge.set_vertex(1, optimizer.vertex(0))  # 相机顶点
            edge.set_measurement(observations[j])  # 观测值 (u, v)
            edge.set_information(np.identity(2))  # 信息矩阵
            edge.set_parameter_id(0, 0)  # 相机参数ID
            optimizer.add_edge(edge)

    # 运行优化
    optimizer.initialize_optimization()
    optimizer.optimize(100)  # 迭代次数

    # 获取优化后的相机参数和三维点位置
    optimized_camera_poses = []
    v_se3 = optimizer.vertex(0)
    pose = v_se3.estimate()
    optimized_camera_poses.append(pose.to_homogeneous_matrix())

    # 计算投影矩阵P
    P = optimized_camera_poses[0][:4, :4]
    return P[:3, :3], P[:3, 3, None]


def test_pose_estimation(observations, points_3d, last_pose=None):
    # 初始误差
    sse = [0, 0]
    K = np.array([[256.204, 0, 256.5], [0, 256.204, 256.5], [0, 0, 1]])
    init_pose = (
        initial_pose(K, np.array(observations), np.array(points_3d))
        if last_pose is None
        else last_pose
    )
    P_ba = bundle_adjustment(K, observations, points_3d, init_pose)
    # 消除二义性
    ## 直接线性变换
    n = init_pose[0] @ np.array([0, 0, 1])
    P_dlt = np.concatenate(init_pose, axis=1) / np.sign(n[2])
    # print("DLT:\n{0}".format(P_dlt))
    uv = np.array(observations).T
    uv = np.pad(uv, ((0, 1), (0, 0)), "constant", constant_values=1)
    XYZ = np.array(points_3d).T
    XYZ = np.pad(XYZ, ((0, 1), (0, 0)), "constant", constant_values=1)
    uv_dlt = K @ P_dlt @ XYZ
    uv_dlt = uv_dlt[:2] / uv_dlt[2]
    sse[0] = np.mean((uv_dlt - uv[:2]) ** 2)

    ## 光束法平差
    n = P_ba[0] @ np.array([0, 0, 1])
    P_ba = np.concatenate(P_ba, axis=1) / np.sign(n[2])
    # print("BA:\n{0}".format(P_ba))
    uv_ba = K @ P_ba @ XYZ
    uv_ba = uv_ba[:2] / uv_ba[2]
    sse[1] = np.mean((uv_ba - uv[:2]) ** 2)

    # print("\nRMSE:")
    # print("before optimization:", np.sqrt(sse[0]))
    # print("after  optimization:", np.sqrt(sse[1]))

    # 提取位置和姿态
    position = P_ba[:3, 3]  # 提取平移部分
    rotation_matrix = P_ba[:3, :3]  # 提取旋转部分

    # print("Position:\n", -rotation_matrix.T @ position)
    # print(
    #     "Orientation (Quaternion):\n", Rotation.from_matrix(rotation_matrix).as_quat()
    # )
    return (
        sse,
        (
            -rotation_matrix.T @ position,
            Rotation.from_matrix(rotation_matrix).as_quat(),
        ),
        P_ba,
    )


def main():
    args = arg_parse()
    # collected_dir = "/disk527/sdb1/a804_cbf/datasets/collect_data/2024-08-19_16-49-01"
    collected_dir = args.collect_dir
    with open(f"{collected_dir}/pose.csv", "r") as f:
        f.readline()  # 去除标题
        flag = False
        with open(f"{collected_dir}/gt_labels.txt", "r") as catelog:
            for i, time in enumerate(map(lambda x: x.strip(), catelog)):
                if float(time) < 22:
                    continue
                observations, points_3d = read_points(
                    os.path.join(collected_dir, "points_id", f"{time}.txt")
                )
                if flag:
                    init_pose = (
                        H[:3, :3],
                        H[:3, 3][:, None],
                    )
                    sse, pose_esti, H = test_pose_estimation(
                        observations, points_3d, init_pose
                    )
                else:
                    sse, pose_esti, H = test_pose_estimation(observations, points_3d)
                for line in f:
                    # 这个判断语句有点问题，如果把阈值改小到5e-4，就会在某个点处产生错误，然后因为不能更新flag，导致后续的所有点都会出错
                    if abs(float(line.split(",")[0]) - float(time)) < 2e-4:
                        pose = [float(x) for x in line.strip().split(",")[1:]]
                        break
                    elif float(line.split(",")[0]) > float(time):
                        pose = [float(x) for x in line.strip().split(",")[1:]]
                        break
                Real_H = real_pose(pose)
                # sse 重投影误差
                # 位姿误差，位置用欧式距离，姿态用四元数距离
                err_dist = np.linalg.norm(pose_esti[0] - np.array(pose[:3]))
                err_rot = (
                    Rotation.from_matrix(Real_H[:3, :3])
                    * Rotation.from_matrix(H[:3, :3]).inv()
                ).as_euler("zyx", degrees=True)

                # print("real position:\n{0}".format(-Real_H[:3, :3].T @ Real_H[:3, 3]))
                # print(
                #     "real orientation:\n{0}".format(
                #         Rotation.from_matrix(Real_H[:3, :3]).as_quat()
                #     )
                # )
                print(
                    "Time: {0}, SSE:{1}, DIST: {2}, ROT: {3}".format(
                        time, sse, err_dist, err_rot
                    )
                )
